import random
import numpy as np
import tensorflow as tf

#作用：定义一个函数 MNIST_DataLabel，用于将标签转换为独热编码（One-Hot Encoding）。
#首先创建一个 10x10 的零矩阵 labels，然后将对角线元素设为 1，表示每个标签的独热编码形式。
def MNIST_DataLabel(data,label):
  labels = np.zeros([10,10])
  for i in range(0,10):
      labels[i][i] = 1
#作用：将数据按标签分组。创建一个包含 10 个空列表的列表 labeled_data，用于存储每个类别的样本。
#遍历所有标签，将对应的数据添加到相应的类别列表中。
  labeled_data = [[],[],[],[],[],[],[],[],[],[]]
  for i in range(0,len(label)):
    label_int = int(label[i])
    labeled_data[label_int].append(data[i])
#作用：对每个类别的样本进行随机打乱，以增加数据的随机性和多样性。
  for i in range(0,10):
    random.shuffle(labeled_data[i])
#作用：返回按标签分组的 labeled_data 和独热编码的 labels。
  return labeled_data,labels
#作用：定义一个函数 Root_Dataset，用于生成根服务器的数据集。
#初始化 root_data 和 root_label 列表，以及一个长度为 10 的零数组 idx，用于记录每个类别已选取的样本数量
def Root_Dataset(labeled_data,labels,root_dataset_size,root_dataset_bias):
  root_data = []
  root_label = []
  idx = np.zeros([10],int)
#作用：从每个类别中按比例选取样本构建根服务器数据集。
#根据 root_dataset_bias 的值，决定是否选择标签为 0 的样本。
#将选取的样本添加到 root_data 和 root_label 中，并更新 idx。
  for i in range(0,root_dataset_size):
    rand = random.random()
    if rand < root_dataset_bias:
      label = 0
    else:
      label = random.randint(0,8)+1
    root_data.append(labeled_data[label][idx[label]])
    idx[label] += 1
    root_label.append(labels[label])
#作用：返回根服务器的数据集 root_data、标签 root_label 和每个类别的样本计数 idx。
  return root_data,root_label,idx

#作用：定义一个函数 Choose_client，用于根据给定的参数选择客户端。
# 计算每个类别的客户端数量 group_size，根据概率 q 决定客户端的标签 label，然后随机选择该标签下的一个客户端。
def Choose_client(l, q, nclass, nclient):
  group_size = nclient/nclass
  rand = random.random()
  if rand < q:
    label = l
  else:
    label = random.randint(0,8)
    if label >= l:
      label += 1
  group_idx = random.randint(0,group_size-1)
  client = label*group_size+group_idx
  # print(label,group_idx)
  return int(client)


#作用：定义一个函数 Load_MNIST，用于加载和预处理 MNIST 数据集。
#使用 TensorFlow 的 load_data 方法加载数据，将像素值缩放到 [0, 1] 范围。
def Load_MNIST(q, nclient = 100, root_dataset_size=100, root_dataset_bias=0.1):
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
  x_train, x_test = x_train / 255.0, x_test / 255.0
#作用：调用之前定义的函数，将训练数据 x_train 和标签 y_train 转换为独热编码，并生成根服务器的数据集。
  labeled_data,labels = MNIST_DataLabel(x_train, y_train)
  root_data,root_label,idx = Root_Dataset(labeled_data,labels,root_dataset_size,root_dataset_bias)
#作用：初始化客户端的数据和标签存储。
#创建 clients_data 和 clients_label 列表，每个包含 nclient 个空列表，用于存储每个客户端的样本数据和标签。
  nclass = 10
  clients_data = []
  clients_label = []
  for i in range(0,nclient):
    clients_data.append([]);
    clients_label.append([]);
#作用：将数据分配给客户端。遍历每个类别的样本，
#根据 Choose_client 函数确定将样本分配给哪个客户端，并将样本数据和标签添加到相应的客户端列表中。
  for i in range(0,nclass):
    for j in range(idx[i],len(labeled_data[i])):
      client_id = Choose_client(i,q,nclass,nclient)
      clients_data[client_id].append(labeled_data[i][j])
      clients_label[client_id].append(labels[i])
#作用：将每个客户端的数据和标签转换为 NumPy 数组，并调整数据的形状，以符合模型输入的要求。
#具体来说，将数据从二维数组（每个样本为 28x28）转换为三维数组（28x28x1），表示灰度图像的高度、宽度和通道数。
  root_data = np.array(root_data)
  root_data = root_data.reshape((root_data.shape[0], 28, 28, 1))
  root_label = np.array(root_label)
  # root_label = root_label.reshape((root_label.shape[0], 10, 1))
  for i in range(0,nclient):
    clients_data[i] = np.array(clients_data[i])
    clients_data[i] = clients_data[i].reshape((clients_data[i].shape[0], 28, 28, 1))
    clients_label[i] = np.array(clients_label[i])
    # clients_label[i] = clients_label[i].reshape((clients_label[i].shape[0],10,1))
#作用：将测试标签 y_test 转换为独热编码格式，以便于分类任务的处理。
#然后，将测试数据 x_test 转换为 NumPy 数组，并调整其形状，添加一个维度表示通道数（灰度图像的通道数为 1）。
  y_test = tf.keras.utils.to_categorical(y_test)
  # y_test = y_test.reshape((y_test.shape[0],10,1))
  x_test = np.array(x_test)
  x_test = x_test.reshape((x_test.shape[0],28,28,1))
#作用：返回处理后的数据，包括根服务器的数据和标签、每个客户端的数据和标签，以及测试数据和标签。这些数据将用于后续的模型训练和评估。
  return root_data,root_label,clients_data,clients_label,x_test,y_test

def Load_FashionMNIST(q, nclient = 100, root_dataset_size=100, root_dataset_bias=0.1):
  # 使用 TensorFlow 的 load_data 方法加载 FashionMNIST 数据集
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()
  # 将像素值缩放到 [0, 1] 范围
  x_train, x_test = x_train / 255.0, x_test / 255.0
  # 调用 MNIST_DataLabel 函数，将训练数据和标签转换为独热编码，并按标签分组
  labeled_data,labels = MNIST_DataLabel(x_train, y_train)
  # 调用 Root_Dataset 函数，生成根服务器的数据集
  root_data,root_label,idx = Root_Dataset(labeled_data,labels,root_dataset_size,root_dataset_bias)
  # 初始化客户端的数据和标签存储
  nclass = 10
  clients_data = []
  clients_label = []
  for i in range(0,nclient):
    clients_data.append([])
    clients_label.append([])
  # 将数据分配给客户端
  for i in range(0,nclass):
    for j in range(idx[i],len(labeled_data[i])):
      # 根据 Choose_client 函数确定将样本分配给哪个客户端
      client_id = Choose_client(i,q,nclass,nclient)
      # 将样本数据添加到相应的客户端列表中
      clients_data[client_id].append(labeled_data[i][j])
      # 将样本标签添加到相应的客户端列表中
      clients_label[client_id].append(labels[i])
  # 将根服务器的数据和标签转换为 NumPy 数组，并调整数据的形状
  root_data = np.array(root_data)
  root_data = root_data.reshape((root_data.shape[0], 28, 28, 1))
  root_label = np.array(root_label)
  # 将每个客户端的数据和标签转换为 NumPy 数组，并调整数据的形状
  for i in range(0,nclient):
    clients_data[i] = np.array(clients_data[i])
    clients_data[i] = clients_data[i].reshape((clients_data[i].shape[0], 28, 28, 1))
    clients_label[i] = np.array(clients_label[i])
    # clients_label[i] = clients_label[i].reshape((clients_label[i].shape[0],10,1))
  # 将测试标签转换为独热编码格式
  y_test = tf.keras.utils.to_categorical(y_test)
  # y_test = y_test.reshape((y_test.shape[0],10,1))
  # 将测试数据转换为 NumPy 数组，并调整其形状
  x_test = np.array(x_test)
  x_test = x_test.reshape((x_test.shape[0],28,28,1))
  # 返回处理后的数据，包括根服务器的数据和标签、每个客户端的数据和标签，以及测试数据和标签
  return root_data,root_label,clients_data,clients_label,x_test,y_test